"""
Standardized output types for consistent model outputs.

This module defines standardized dataclasses for model outputs to ensure
consistent return types across different components of the framework.
"""

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

import torch

EMPTY_TENSOR = torch.zeros(0)


@dataclass
class HookOutput:
    """
    Standardized output from hook operations.

    This ensures consistent return type structure across different hooks.

    Args:
        output: The primary output tensor or dictionary
        modified: Whether the output was modified by the hook
        extra_data: Additional data that doesn't fit into the primary output
    """

    output: Union[torch.Tensor, Dict[str, Any]]
    modified: bool = False
    extra_data: Optional[Dict[str, Any]] = None


@dataclass
class ModelOutput:
    """
    Base class for all model outputs.

    Similar to Hugging Face's implementation, this standardizes model outputs
    to make interfaces more consistent and easier to work with.

    Attributes:
        features: Feature representation extracted by the model
        logits: Classification logits (unnormalized predictions)
        hidden_states: All intermediate hidden states (for certain architectures)
        attentions: Attention weights (for transformer-based models)
        extra_outputs: Additional outputs that don't fit into standard fields
    """

    features: torch.Tensor = field(default_factory=lambda: EMPTY_TENSOR)
    logits: torch.Tensor = field(default_factory=lambda: EMPTY_TENSOR)
    hidden_states: Optional[tuple[torch.Tensor, ...]] = None
    attentions: Optional[tuple[torch.Tensor, ...]] = None
    extra_outputs: Optional[Dict[str, Any]] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert the dataclass to a dictionary, excluding None values."""
        result = {k: v for k, v in self.__dict__.items() if v is not None}
        extra_outputs = result.pop("extra_outputs", None)
        if extra_outputs:
            result.update(extra_outputs)
        return result

    def __getitem__(self, key):
        """Enable dictionary-like access to attributes."""
        return getattr(self, key)


@dataclass
class AdapterOutput(ModelOutput):
    """
    Standardized output from adapter modules.

    This ensures consistent structure when adapters return features and additional states.

    Args:
        features: The primary feature tensor
        adapter_hidden_states: Optional hidden states from the adapter
    """

    adapter_hidden_states: Optional[List[torch.Tensor]] = None


@dataclass
class BackboneOutput(ModelOutput):
    """
    Specialized output type for backbone models.

    Contains backbone-specific outputs like extracted features,
    while maintaining compatibility with the base ModelOutput.

    Attributes:
        pooled_output: Pooled feature representation (e.g., CLS token)
        last_hidden_state: The last hidden state of the model
    """

    pooled_output: torch.Tensor = field(default_factory=lambda: EMPTY_TENSOR)
    last_hidden_state: torch.Tensor = field(default_factory=lambda: EMPTY_TENSOR)
